import json
from pymongo import MongoClient
from bson import ObjectId
from datetime import datetime

class MongoDBExporter:
    def __init__(self):
        self.client = self._get_mongodb_connection()
    
    def _get_mongodb_connection(self):
        """获取MongoDB连接"""
        username = "root"
        password = "pj9gbcht"
        host = "dbconn.sealosbja.site"
        port = 31140
        
        # 构建连接URI
        uri = f"mongodb://{username}:{password}@{host}:{port}/?directConnection=true"
        return MongoClient(uri)
    
    def _convert_to_json_serializable(self, obj):
        """递归处理所有数据类型，确保JSON可序列化"""
        if isinstance(obj, ObjectId):
            return str(obj)
        elif isinstance(obj, datetime):
            return obj.isoformat()  # 转换为ISO 8601格式字符串
        elif isinstance(obj, list):
            return [self._convert_to_json_serializable(item) for item in obj]
        elif isinstance(obj, dict):
            return {key: self._convert_to_json_serializable(value) for key, value in obj.items()}
        else:
            # 处理其他非标准类型
            try:
                json.dumps(obj)
                return obj
            except TypeError:
                return str(obj)
    
    # def export_data(self, db_name="AI25", collection_name="AI25LIST", output_file="AI25LIST.json"):
    # def export_data(self, db_name="AI25", collection_name="loginuser", output_file="loginuser.json"):
    # def export_data(self, db_name="AI25", collection_name="POOL", output_file="POOL.json"):
    def export_data(self, db_name="AI25", collection_name="TR35", output_file="TR35.json"):
        """
        导出MongoDB集合数据到JSON文件
        
        参数:
        db_name: 数据库名称 (默认: 'AI25')
        collection_name: 集合名称 (默认: 'AI25LIST')
        output_file: 输出文件名 (默认: 'AI25LIST.json')
        """
        # 选择数据库和集合
        db = self.client[db_name]
        collection = db[collection_name]
        
        # 查询所有数据
        data = list(collection.find({}))
        
        # 递归转换所有特殊类型
        converted_data = []
        for doc in data:
            # 显式处理您的特殊字段
            if 'processed_at' in doc:
                doc['processed_at'] = self._convert_to_json_serializable(doc['processed_at'])
            if 'created_at' in doc:
                doc['created_at'] = self._convert_to_json_serializable(doc['created_at'])
            if 'updated_at' in doc:
                doc['updated_at'] = self._convert_to_json_serializable(doc['updated_at'])
            
            # 递归处理整个文档
            converted_data.append(self._convert_to_json_serializable(doc))
        
        # 保存为JSON文件
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(converted_data, f, ensure_ascii=False, indent=2)
        
        return len(data)  # 返回导出的记录数

# 使用示例
if __name__ == "__main__":
    exporter = MongoDBExporter()
    record_count = exporter.export_data()
    print(f"导出完成，共导出 {record_count} 条数据。")